{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pathlib \n",
    "import pickle as pkl\n",
    "import sys\n",
    "sys.path.insert(0, \"/home/estengel/real_good_robot\")\n",
    "from annotate_data import get_pairs, annotate_pairs, flip_pair, rotate_pair\n",
    "\n",
    "\n",
    "# data_home=pathlib.Path(\"/srv/local1/estengel/gr_data//2020-06-05-18-28-46_Real-Stack-SPOT-Trial-Reward-Masked-Testing\")\n",
    "# data_home=pathlib.Path(\"/srv/local1/estengel/gr_data/2020-06-07-16-31-33_Real-Stack-Two-Step-Reward-Masked-Testing\")\n",
    "# data_home=pathlib.Path(\"/srv/local1/estengel/gr_data/2020-06-07-18-10-48_Real-Rows-SPOT-Trial-Reward-Masked-Testing\")\n",
    "# data_home=pathlib.Path(\"/srv/local1/estengel/gr_data//2020-06-07-17-19-34_Real-Rows-Two-Step-Reward-Masked-Testing/\")\n",
    "\n",
    "\n",
    "# data_home = pathlib.Path(\"/srv/local1/estengel/gr_data_sim/2021-02-13-16-43-06_Sim-Rows-Two-Step-Reward-Masked-Testing-Three-Step-History\")\n",
    "\n",
    "\n",
    "\n",
    "data_home = pathlib.Path(\"/srv/local1/estengel/gr_data_sim/2021-02-13-16-30-37_Sim-Stack-SPOT-Trial-Reward-Masked-Testing-Three-Step-History/\") \n",
    "data_output = pathlib.Path(\"/srv/local1/estengel/gr_data_sim_long_command/2021-02-13-16-30-37_Sim-Stack-SPOT-Trial-Reward-Masked-Testing-Three-Step-History/\") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "outputPrepend"
    ]
   },
   "outputs": [],
   "source": [
    "# pairs_with_actions = get_pairs(data_home, is_sim=True)\n",
    "pairs_with_actions = get_pairs(data_home, is_sim=True, is_row = False, w=20, filter_colors=False, long_command = True)\n",
    "print(len(pairs_with_actions))\n",
    "# pairs_with_actions = annotate_pairs(pairs, is_stack=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "with open(data_output.joinpath(\"with_actions.pkl\"),\"wb\") as f1:\n",
    "    pkl.dump(pairs_with_actions, f1)\n",
    "\n",
    "# with open(data_home.joinpath(\"with_actions.pkl\"), \"rb\") as f1:\n",
    "#     pairs_with_actions = pkl.load(f1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np \n",
    "from matplotlib import pyplot as plt\n",
    "import cv2\n",
    "np.random.seed(10)\n",
    "idxs = np.random.randint(0, len(pairs_with_actions)-1, 30)\n",
    "# idxs = [241, 418, 373, 89, 417, 352]\n",
    "# for idx in idxs:\n",
    "#     test_p = pairs_with_actions[idx]\n",
    "#     test_p.show()\n",
    "#     print(idx)\n",
    "#     print(test_p.generate())\n",
    "print(len(pairs_with_actions))\n",
    "test_p = pairs_with_actions[33]\n",
    "\n",
    "# prev_depthmap = test_p.prev_image[:,:,3:]\n",
    "\n",
    "\n",
    "# plt.imshow(prev_depthmap)\n",
    "# plt.figure()\n",
    "# np.max(prev_depthmap)\n",
    "test_p.show()\n",
    "test_p.generate()\n",
    "\n",
    "#     # next = input(\"next: \")\n",
    "# test_p_1 = flip_pair(test_p, axis=1)\n",
    "# test_p_1.show()\n",
    "# test_p_1.generate()\n",
    "\n",
    "# test_p_1 = rotate_pair(test_p, 1)\n",
    "# test_p_1.show()\n",
    "# test_p_1.generate()\n",
    "\n",
    "# test_p_2 = flip_pair(test_p, axis=2)\n",
    "# test_p_2.show()\n",
    "# test_p_2.generate()\n",
    "\n",
    "# test_p_3 = flip_pair(test_p, axis=3)\n",
    "# test_p_3.show()\n",
    "# test_p_3.generate()\n",
    "# test_p_4 = flip_pair(test_p, axis=4)\n",
    "# test_p_4.show()\n",
    "# test_p_4.generate()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(test_p_1.next_image.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "python361264bitblockscondae1e5e96102154286a7dee4f769994df8",
   "display_name": "Python 3.6.12 64-bit ('blocks': conda)"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.12"
  },
  "metadata": {
   "interpreter": {
    "hash": "80b57d3b473cded757c1c09da496d17d13901a1bde76b141e0b4105d5712fb37"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}